"""
求解器模块

该模块实现了两相流守恒方程的求解，包括：
1. 质量守恒方程
2. 动量守恒方程
3. 能量守恒方程
4. 状态方程
"""

import numpy as np
from typing import Dict, Optional, Tuple
from .staggered_grid import StaggeredGrid
from .numerical import NumericalMethods
from ..physics.two_phase import TwoPhaseFlow

class ConservationSolver:
    """守恒方程求解器类"""
    
    def __init__(
        self,
        grid: StaggeredGrid,
        numerical: NumericalMethods,
                 two_phase: TwoPhaseFlow,
        config: dict
    ):
        """
        初始化求解器
        
        参数:
            grid (StaggeredGrid): 交错网格对象
            numerical (NumericalMethods): 数值方法对象
            two_phase (TwoPhaseFlow): 两相流模型对象
            config (dict): 配置参数字典
        """
        self.grid = grid
        self.numerical = numerical
        self.two_phase = two_phase
        self.config = config
        
        # 获取求解器参数
        self.max_iterations = config['numerical']['max_iterations']
        self.convergence_criteria = config['numerical']['convergence_criteria']
        self.relaxation_factors = config['numerical']['relaxation_factors']
        
    def solve_time_step(self, dt: float) -> bool:
        """
        求解一个时间步
        
        参数:
            dt (float): 时间步长
            
        返回:
            bool: 是否收敛
        """
        # 保存当前时间步的物理量
        self.grid.store_current_step()
        
        # SIMPLE算法迭代
        for iter in range(self.max_iterations):
            # 计算动量方程系数矩阵
            A_m_x, b_m_x = self.numerical._build_axial_momentum_matrix(self.grid)
            A_m_y, b_m_y = self.numerical._build_lateral_momentum_matrix(self.grid)
            
            # 求解动量方程
            self._solve_momentum_equations(A_m_x, A_m_y, b_m_x, b_m_y)
            
            # 计算压力修正方程系数矩阵
            A_p = self.numerical.build_pressure_correction_matrix(self.grid)
            
            # 求解压力修正方程
            self._solve_pressure_correction_equation(A_p)
            
            # 更新速度场
            self.numerical.correct_velocities(self.grid, self.relaxation_factors['velocity'])
            
            # 更新压力场
            self.numerical.correct_pressure(self.grid, self.relaxation_factors['pressure'])
            
            # 求解能量方程
            self._solve_energy_equation()
            
            # 计算相变源项
            source_terms = self.two_phase.calculate_phase_change()
            
            # 更新物理量
            self._update_variables(source_terms, dt)
            
            # 检查收敛性
            residuals = self.numerical.calculate_residuals(self.grid, A_m_x, A_m_y, A_p)
            
            if self._check_convergence(residuals):
                return True
                
        return False
        
    def _solve_momentum_equations(self, A_m: Dict[str, np.ndarray]):
        """
        求解动量方程
        
        参数:
            A_m (Dict[str, np.ndarray]): 动量方程系数矩阵
        """
        # 计算界面力和壁面摩擦力
        F_i = self.two_phase.calculate_interfacial_forces()
        F_w = self.two_phase.calculate_wall_friction()
        
        # 求解轴向动量方程
        b_x = self._build_momentum_source_term(
            'x',
            F_i['x'],
            F_w['x']
        )
        u = self.numerical.solve_linear_system(
            A_m['x'],
            b_x,
            self.grid.get_axial_velocity_field()
        )
        self.grid.set_axial_velocity_field(u)
        
        # 求解横向动量方程
        b_y = self._build_momentum_source_term(
            'y',
            F_i['y'],
            F_w['y']
        )
        v = self.numerical.solve_linear_system(
            A_m['y'],
            b_y,
            self.grid.get_lateral_velocity_field()
        )
        self.grid.set_lateral_velocity_field(v)
        
    def _solve_pressure_correction_equation(self, A_p: np.ndarray):
        """
        求解压力修正方程
        
        参数:
            A_p (np.ndarray): 压力修正方程系数矩阵
        """
        # 计算质量不平衡
        b = self.numerical.calculate_mass_imbalance(self.grid)
        
        # 求解压力修正方程
        p_prime = self.numerical.solve_linear_system(
            A_p,
            b,
            np.zeros_like(b)
        )
        
        # 更新压力场
        self.grid.set_scalar_field(
            'pressure_correction',
            p_prime
        )
        
    def _solve_energy_equation(self):
        """求解能量方程"""
        # 计算传热量
        q = self.two_phase.calculate_heat_transfer()
        
        # 构建能量方程系数矩阵
        A_e = self.numerical.build_energy_matrix(
            self.grid,
            self.two_phase
        )
        
        # 构建源项
        b = self._build_energy_source_term(q)
        
        # 求解能量方程
        h = self.numerical.solve_linear_system(
            A_e,
            b,
            self.grid.get_scalar_field('enthalpy')
        )
        
        # 更新焓场
        self.grid.set_scalar_field('enthalpy', h)
        
    def _build_momentum_source_term(
        self,
        direction: str,
        F_i: np.ndarray,
        F_w: np.ndarray
    ) -> np.ndarray:
        """
        构建动量方程源项
        
        参数:
            direction (str): 方向 ('x' 或 'y')
            F_i (np.ndarray): 界面力
            F_w (np.ndarray): 壁面摩擦力
            
        返回:
            np.ndarray: 源项向量
        """
        # 获取物理量
        rho = self.grid.get_scalar_field('density')
        p = self.grid.get_scalar_field('pressure')
        
        # 计算压力梯度
        if direction == 'x':
            dp_dx = np.gradient(p, self.grid.dx, axis=0)
            g_term = rho * 9.81  # 重力项
        else:
            dp_dy = np.gradient(p, self.grid.dy, axis=1)
            g_term = 0.0  # 忽略横向重力
            
        # 构建源项
        b = -(F_i + F_w) - g_term
        
        return b
        
    def _build_energy_source_term(
        self,
        q: Dict[str, np.ndarray]
    ) -> np.ndarray:
        """
        构建能量方程源项
        
        参数:
            q (Dict[str, np.ndarray]): 传热量字典
            
        返回:
            np.ndarray: 源项向量
        """
        # 计算总传热量
        q_total = q['convection'] + q['boiling'] + q['wall']
        
        return q_total
        
    def _update_variables(
        self,
        source_terms: Dict[str, np.ndarray],
        dt: float
    ):
        """
        更新物理量
        
        参数:
            source_terms (Dict[str, np.ndarray]): 相变源项字典
            dt (float): 时间步长
        """
        # 更新密度
        rho = self.grid.get_scalar_field('density')
        rho += source_terms['mass'] * dt
        self.grid.set_scalar_field('density', rho)
        
        # 更新速度
        u = self.grid.get_axial_velocity_field()
        v = self.grid.get_lateral_velocity_field()
        u += source_terms['momentum_x'] * dt / rho
        v += source_terms['momentum_y'] * dt / rho
        self.grid.set_axial_velocity_field(u)
        self.grid.set_lateral_velocity_field(v)
        
        # 更新焓
        h = self.grid.get_scalar_field('enthalpy')
        h += source_terms['energy'] * dt / rho
        self.grid.set_scalar_field('enthalpy', h)
        
        # 更新温度和空泡率
        self.grid.calculate_temperature()
        self.grid.calculate_void_fraction()
        
    def _check_convergence(
        self,
        residuals: Dict[str, float]
    ) -> bool:
        """
        检查收敛性
        
        参数:
            residuals (Dict[str, float]): 残差字典
            
        返回:
            bool: 是否收敛
        """
        # 检查各个方程的残差
        for name, value in residuals.items():
            if value > self.convergence_criteria[name]:
                return False
                
        return True